#!/usr/bin/env python2

import sys
import os
import base64
import math
import zipfile
import datetime



def decode_txt_to_zip(txt_path):
    try:
        # Process a single file
        if '_part' not in os.path.basename(txt_path):
            with open(txt_path, 'r') as f:
                encoded = f.read().strip()
        else:
            # Process split files
            base_path = txt_path.split('_part')[0]
            parts = []
            i = 1
            while True:
                part_path = "{}_part{:03d}.txt".format(base_path, i)
                if not os.path.exists(part_path):
                    break
                with open(part_path, 'r') as f:
                    parts.append(f.read().strip())
                i += 1

            if not parts:
                raise ValueError("No split files found")

            encoded = ''.join(parts)

        decoded = base64.b64decode(encoded)
        zip_path = "{}.zip".format(os.path.splitext(base_path if '_part' in txt_path else txt_path)[0])
        with open(zip_path, 'wb') as f:
            f.write(decoded)
        print("Successfully decoded to: {}".format(zip_path))
        return zip_path

    except Exception as e:
        print("Decoding failed:", str(e))
        sys.exit(1)

def extract_zip_to_folder(zip_path):
    try:
        # Generate folder name with timestamp
        base_name = os.path.splitext(zip_path)[0]
        extract_folder = "{}".format(base_name)

        # Create target folder
        if not os.path.exists(extract_folder):
            os.makedirs(extract_folder)

        # Extract ZIP file
        with zipfile.ZipFile(zip_path, 'r') as zip_ref:
            zip_ref.extractall(extract_folder)

        print("Successfully extracted to: {}".format(extract_folder))
        return extract_folder
    except Exception as e:
        print("Extraction failed:", str(e))
        sys.exit(1)


def compress_folder_to_zip(folder_path, zip_path):
    try:
        with zipfile.ZipFile(zip_path, 'w', zipfile.ZIP_DEFLATED) as zipf:
            for root, dirs, files in os.walk(folder_path):
                arc_dir = os.path.relpath(root, folder_path)
                if arc_dir != '.':
                    dir_arcname = arc_dir + '/'
                    zip_info = zipfile.ZipInfo(dir_arcname)
                    zip_info.external_attr = (0o40755 & 0xFFFF) << 16  # drwxr-xr-x
                    zipf.writestr(zip_info, '')
                for file in files:
                    file_path = os.path.join(root, file)
                    arcname = os.path.relpath(file_path, folder_path)
                    zipf.write(file_path, arcname)
        return True
    except Exception as e:
        print("Compression failed:", str(e))
        return False


def encode_zip_to_txt(zip_path, max_size_kb=None):
    try:
        with open(zip_path, 'rb') as f:
            data = f.read()
        encoded = base64.b64encode(data)

        if max_size_kb:
            max_size_bytes = max_size_kb * 1024
            if len(encoded) > max_size_bytes:
                parts = int(math.ceil(float(len(encoded)) / max_size_bytes))
                base_name = os.path.splitext(zip_path)[0]
                for i in range(parts):
                    start = i * max_size_bytes
                    end = (i + 1) * max_size_bytes
                    part_data = encoded[start:end]
                    part_path = "{}_part{:03d}.txt".format(base_name, i + 1)
                    with open(part_path, 'w') as f:
                        f.write(part_data)
                print("Successfully split into {} parts (each max {} KB)".format(parts, max_size_kb))
                return

        txt_path = os.path.splitext(zip_path)[0] + '.txt'
        with open(txt_path, 'w') as f:
            f.write(encoded)
        print("Successfully encoded to: {}".format(txt_path))
    except Exception as e:
        print("Encoding failed:", str(e))
        sys.exit(1)





def main():
    # sys.argv contains all arguments passed to the script, where the first argument is the script name itself

    if len(sys.argv) < 2 or len(sys.argv) > 3:
        print "Usage:"
        print "script.py aaa.txt"
        print "script.py folder_path"
        print "script.py folder_path 1024"
        return

    input_path = sys.argv[1]

    if os.path.splitext(input_path)[1].lower() == ".txt":
        zip_path = decode_txt_to_zip(input_path)
        extract_zip_to_folder(zip_path)
        return
    elif os.path.isdir(input_path):
        txt_max_size_KB = None
        if len(sys.argv) == 3:
            # Two arguments provided, assign the second argument to txt_max_size_KB
            try:
                txt_max_size_KB = int(sys.argv[2])
            except ValueError:
                print "The second argument must be an integer"
                return

        timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")

        folder_path = os.path.abspath(input_path)
        folder_name = os.path.basename(folder_path)
        parent_dir = os.path.dirname(folder_path)
        zip_path = os.path.join(parent_dir, "{}_{}.zip".format(folder_name, timestamp))

        if not compress_folder_to_zip(folder_path, zip_path):
            sys.exit(1)
        encode_zip_to_txt(zip_path, txt_max_size_KB)

    else:
        print "First argument is neither txt nor folder"
        return



if __name__ == "__main__":
    main()
